# Copyright (c) 2024, yuxinyong and contributors
# For license information, please see license.txt

import os
import frappe
from frappe.model.document import Document
import json
from pathlib import Path
from openai import OpenAI


class FreightBooking(Document):
    def validate(self):
        self.parse_attachment()

    def parse_attachment(self):
        if self.attachment and not self.content:
            json_data = get_json_data_from_attachment(self.attachment)
            if json_data:
                self.content = str(json_data)
                self.update(json_data)

@frappe.whitelist()
def upload_file(docname , doctype , filename , filedata) :
    base_dir = frappe.get_site_path()
    # 保存文件到 public/invoice_upload 目录
    user = frappe.get_user()
    upload_dir = os.path.join(base_dir , 'public/files/' + doctype , user.name)
    if not os.path.exists(upload_dir) :
        os.makedirs(upload_dir)
    # 分割文件名和扩展名
    file_name , file_extension = os.path.splitext(filename)
    if file_extension not in ['.PDF', '.pdf', '.PNG', '.png', '.JPG', '.jpg', '.doc', '.docx', '.xls', '.xlsx']:
        frappe.throw("仅支持PDF、PNG、JPG、doc, xls 格式的文件上传！")
    #new_filename = f"{docname}{file_extension}"
    upload_path = os.path.join(upload_dir , filename)    
    with open(upload_path , 'wb') as f :
        f.write(filedata)
    doc = frappe.get_doc(doctype , docname)
    doc.file_path = upload_path
    doc.save()

def get_json_data_from_attachment(file_path):
    base_dir = frappe.get_site_path()
    if (file_path[:6] == "/files"):
        file_path = base_dir + "/public" + file_path
    else:
        file_path = base_dir + file_path
    file_path = os.path.join(frappe.utils.get_bench_path(),'sites',file_path[2 if file_path[0]=='.' else 1:])

    settings = frappe.get_cached_doc("Agent Assistant Settings")
    client = get_ai_client(settings)
    file_object = client.files.create(file=Path(file_path), purpose="file-extract")    
    file_content = client.files.content(file_id=file_object.id).text
    prompt = """请按以下要求从文件中提取内容
    输出json格式，
    提取这些key的内容：shipper,consignee,notify_party,loading_port,destination_port, ref_no,mark,description,qty,gross_weight,net_weight, measurement,
    添加一个固定key值 doctype:'Shipping Track'
    shipper, consignee, notify_party 公司名，地址及电话合并成一个文本，不用单独分开
    数量，净，毛重仅输出数字不加引号
    """ 
    # 把文件内容通过系统提示词 system prompt 放进请求中
    messages = [
        {
            "role": "system",
            "content": "你是 Kimi，由 Moonshot AI 提供的人工智能助手，你更擅长中文和英文的对话。你会为用户提供安全，有帮助，准确的回答。同时，你会拒绝一切涉及恐怖主义，种族歧视，黄色暴力等问题的回答。Moonshot AI 为专有名词，不可翻译成其他语言。",
        },
        {
            "role": "system",
            "content": file_content, # <-- 这里，我们将抽取后的文件内容（注意是文件内容，而不是文件 ID）放置在请求中
        },
        {"role": "user", "content": settings.prompt},
    ]
    
    # 然后调用 chat-completion, 获取 Kimi 的回答
    completion = client.chat.completions.create(
        model="moonshot-v1-32k",
        messages=messages,
        temperature=0.3,
    )
    content = completion.choices[0].message.content
    json_str = content.split('```json')[1].split('```')[0].strip()
    data = json.loads(json_str)
    return data

def get_ai_client(settings):    
    ai_api_key = settings.get_password("kimi_api_key")
    client = OpenAI(api_key=ai_api_key, base_url="https://api.moonshot.cn/v1")
    return client
